import os
import re
import fitz
from paddleocr import PaddleOCR

'''
整理中文文件名工具
pdf文件名规则：年份+省份+地区+[学校]+[学期]+[年级]+[考试]+[科目]+[试题]+[&答案]
pdf_path pdf文件所在目录
期中考|中考？？ @ TODO
'''


# 按照规则整理文件名
def arrange_name(directory: str, filename1: str):
    filename2: str = filename1

    # 解析文件名
    meta = parse_title(filename2)
    # 如果有缺失，从文件内容中获取
    if meta['year'] == '' or meta['section'] == '' or meta['grade'] == '' or meta['subject'] == '':
        meta2 = get_pdf_meta(directory, filename1)
        if meta2:
            if meta['subject'] == '' and meta2['subject'] != '':  # 缺少学科，添加学科
                filename2 = meta2['subject'] + filename2
            if meta['grade'] == '' and meta2['grade'] != '':  # 缺少年级，添加年级
                filename2 = meta2['grade'] + filename2
            if meta['section'] == '' and meta2['section'] != '':  # 缺少学段，添加学段
                filename2 = meta2['section'] + filename2
            if meta['year'] == '' and meta2['year'] != '':  # 缺少年份，添加年份
                filename2 = meta2['year'] + filename2

    # 年份放前面
    pattern1 = re.compile(r'20\d{2}(-20\d{2})?([学年度届]*)*([\.]?\d{1,2}[月]?)*', re.I)
    match = pattern1.search(filename2)
    if match:
        if match.start() > 0:
            year = match.group()
            filename2 = year + pattern1.sub('', filename2, 1)

    # 地区添加省份
    pattern2_0 = re.compile(
        r'(东城|西城|朝阳|丰台|石景山|海淀|顺义|通州|大兴|房山|门头沟|昌平|平谷|密云|怀柔|延庆)[区]?', re.I)
    pattern2_1 = re.compile(r'北京[市]?', re.I)
    match = pattern2_0.search(filename2)
    if match:
        area = match.group()
        match = pattern2_1.search(filename2)
        if not match:
            filename2 = pattern2_0.sub('北京' + area, filename2, 1)
            # print(filename1, '+北京', filename2)

    return filename2


'''
2001-2003学年度
第一学期|上|下
xxxx学校|中学|中
小学|初中|高中|
[一二三四五六七八九]年级，
期中|期末|月考|[一二三]模拟|中考|高考|诊断
语文、数学、英语、历史、地理、政治、物理、化学、生物、美术、音乐、信息与技术、体育、思想品德、道德与法治
试卷、试题、练习、考、答案、真题、模拟、问卷
'''


# 打开pdf内容并获取年份，学科，年级等关键字
def get_pdf_meta(directory: str, filename1: str):
    # 先打开 PDF 文件
    text = read_pdf_page0(directory, filename1)

    # 如果text为None，开始ocr识别
    if not text:
        text = ocr_pdf_page0(directory, filename1)

    # 解析标题
    if text:
        meta = parse_title(text)
        meta['title'] = text[:400]
        return meta

    return None


# 读取pdf内容
def read_pdf_page0(directory: str, filename1: str):
    # 打开 PDF 文件
    with fitz.open(os.path.join(directory, filename1)) as pdf:
        page = pdf.load_page(0)  # pdf[0]
        text = page.get_text()
        if text:
            # 删除空行
            text = '\n'.join([line for line in text.split('\n') if line.strip() != ''])

            # 清除字符串中的空格和空白字符，但保留换行符
            text = re.sub(r'\s+', '', text)

            # 清除分行标记
            text = re.sub(r'\r\n+', '', text)

            # 去除矢量图形标识
            text = re.sub(r'\{#\{.*\}#\}', '', text)
            return text

    return None


# ocr识别pdf内容
def ocr_pdf_page0(directory: str, filename1: str):
    text: str = ''
    ocr = PaddleOCR(use_angle_cls=True, lang="ch", page_num=1, use_gpu=False)
    result = ocr.ocr(os.path.join(directory, filename1), cls=True)
    res = result[0]
    for box, [txt, score] in res:
        text = text + txt
    return text


# 解析pdf标题
def parse_title(filename: str):
    meta = {}
    year: str = ''  # 年份
    section: str = ''  # 学段
    grade: str = ''  # 年级
    test = ''  # 考试类型
    subject: str = ''  # 学科

    # 年份放前面
    pattern_year = re.compile(r'20\d{2}(-20\d{2})?([学年度届]*)*([\.]?\d{1,2}[月]?)*', re.I)
    match = pattern_year.search(filename)
    if match:
        year = match.group()

    # 获取年级 '中考'不能匹配'期中考试'
    pattern1 = re.compile(r'(([^期]中考|高考)|[小初高][学中]?)(([一二三四五六七八九]|十一|十二|十)?(年级)*)', re.I)
    for match in pattern1.finditer(filename):
        # 学段
        sec = match.group(1)
        # 年级
        gra = match.group(3)

        if '中考' in sec:
            test = '中考'
            section = '初中'
            grade = '三年级'
        elif '高考' in sec:
            test = '高考'
            section = '高中'
            grade = '三年级'
        elif '小学' in sec:
            section = '小学'
        elif '初中' in sec:
            section = '初中'
        elif '高中' in sec:
            section = '高中'
        elif '小' in sec and len(gra) > 0:
            section = '小学'
        elif '初' in sec and len(gra) > 0:
            section = '初中'
        elif '高' in sec and len(gra) > 0:
            section = '高中'

        if gra:
            if '一' in gra:
                grade = '一年级'
            elif '二' in gra:
                grade = '二年级'
            elif '三' in gra:
                grade = '三年级'
            elif '四' in gra:
                section = '小学'
                grade = '四年级'
            elif '五' in gra:
                section = '小学'
                grade = '五年级'
            elif '六' in gra:
                section = '小学'
                grade = '六年级'
            elif '七年级' in gra:
                section = '初中'
                grade = '一年级'
            elif '八年级' in gra:
                section = '初中'
                grade = '二年级'
            elif '九年级' in gra:
                section = '初中'
                grade = '三年级'
            elif '十年级' in gra:
                section = '高中'
                grade = '一年级'
            elif '十一年级' in gra:
                section = '高中'
                grade = '二年级'
            elif '十二年级' in gra:
                section = '高中'
                grade = '三年级'
            elif '营' in sec:
                section = '高中'
                grade = '三年级'

    # 学段没有时，可以通过年级获取来修正
    if section == '':
        pattern10 = re.compile(r'([四五六七八九]|十一|十二|十)?年级', re.I)
        match = pattern10.search(filename)
        if match:
            gra = match.group()
            if '四' in gra:
                section = '小学'
                grade = '四年级'
            elif '五' in gra:
                section = '小学'
                grade = '五年级'
            elif '六' in gra:
                section = '小学'
                grade = '六年级'
            elif '七' in gra:
                section = '初中'
                grade = '一年级'
            elif '八' in gra:
                section = '初中'
                grade = '二年级'
            elif '九' in gra:
                section = '初中'
                grade = '三年级'
            elif '十年级' in gra:
                section = '高中'
                grade = '一年级'
            elif '十一年级' in gra:
                section = '高中'
                grade = '二年级'
            elif '十二年级' in gra:
                section = '高中'
                grade = '三年级'

    # 深度修正高考
    if '高等学校招生' in filename:
        test = '高考'
        section = '高中'
        grade = '三年级'

    # 深度修正年级,匹配中考高考栏目
    if grade == '' or '三' in grade:
        pattern11 = re.compile(r'[零一二三][次轮]?[模诊检]|联考|联合|会考|联合|联盟|模拟', re.I)
        match = pattern11.search(filename)
        if match:
            if section == '初中':
                test = '中考'
                if grade == '':
                    grade = '三年级'
            elif section == '高中':
                test = '高考'
                if grade == '':
                    grade = '三年级'

    # 获取学科
    pattern2 = re.compile(r'语文|数学|英语|历史|地理|政治|物理|化学|生物|美术|音乐|信息与技术|体育|思想品德|道德与法治|道法|文综|理综', re.I)
    match = pattern2.search(filename)
    if match:
        subject = match.group()
        subject = re.sub('思想品德|道德与法治|道法', '政治', subject)

    # 修正中考高考栏目
    if subject == '理综' or subject == '文综':
        if test == '' and section == '初中':
            test = '中考'
        elif test == '' and section == '高中':
            test = '高考'

    # 深度修正学段，匹配高考竞赛栏目
    if section == '' or grade == '':
        pattern11 = re.compile(r'竞赛|决赛|复赛|保送|奥林匹克|THUSSAT|科学营|金秋营|综合营|强基|攀登计划|人才培养|CSP|学术能力诊断|国家集训队|标准学术能力|学术能力测试|领军计划|NOI|IChO|TDA|少年班|CSP|全国中学生|大学自主招生', re.I)
        match = pattern11.search(filename)
        if match and section != '初中':
            test = '竞赛'
            section = '高中'

    return {'year': year, 'test': test, 'grade': grade, 'section': section, 'subject': subject}


if __name__ == '__main__':

    pdf_path = 'D:\\work\\文库资源\\03'  # 已经下载的pdf文件目录

    # 遍历文件夹目录里的所有pdf文件
    for root, dirs, files in os.walk(pdf_path):
        for filename in files:
            if filename.endswith('.pdf'):
                new_filename = arrange_name(root, filename)  # 整理文件名
                # print(filename, '->', new_filename)
                if new_filename != filename:  # 如果有改动
                    new_file = os.path.join(root, new_filename)
                    if not os.path.exists(new_file):
                        # pdf文件重命名
                        try:
                            print(filename, '->', new_filename)
                            os.rename(os.path.join(root, filename), new_file)
                        except Exception as e:
                            print(e)
